// NOLINTBEGIN
#pragma once
#include <fmt/base.h>

#include <cstdint>
#include <memory>
#include <string>
#include <unordered_map>
#include <variant>
#include <vector>

#include "task_set.hpp"

////////////////////////////////////////////////////////
//                        PLAN                        //
////////////////////////////////////////////////////////

namespace hashjoin {
using Data = std::variant<int32_t, int64_t, float, double, std::string,
                          std::shared_ptr<std::vector<size_t>>, std::monostate>;
using Table = std::vector<std::vector<Data>>;

void print_table(const Table& data) {
  auto escape_string = [](const std::string& s) {
    std::string escaped;
    for (char c : s) {
      switch (c) {
        case '"':
          escaped += "\\\"";
          break;
        case '\\':
          escaped += "\\\\";
          break;
        case '\n':
          escaped += "\\n";
          break;
        case '\r':
          escaped += "\\r";
          break;
        case '\t':
          escaped += "\\t";
          break;
        default:
          escaped += c;
          break;
      }
    }
    return escaped;
  };

  for (const auto& record : data) {
    std::vector<std::string> parts;
    parts.reserve(record.size());

    for (const auto& field : record) {
      std::string part = std::visit(
          [&escape_string](const auto& arg) {
            using T = std::decay_t<decltype(arg)>;
            using namespace std::string_literals;

            if constexpr (std::is_same_v<T, std::monostate>) {
              return "NULL"s;
            } else if constexpr (std::is_same_v<T, int32_t> ||
                                 std::is_same_v<T, int64_t> ||
                                 std::is_same_v<T, float> ||
                                 std::is_same_v<T, double>) {
              return fmt::format("{}", arg);
            } else if constexpr (std::is_same_v<T, std::string>) {
              return fmt::format("\"{}\"", escape_string(arg));
            } else if constexpr (std::is_same_v<
                                     T, std::shared_ptr<std::vector<size_t>>>) {
              std::string vector_str = "(";
              for (size_t i = 0; i < arg->size(); ++i) {
                vector_str += fmt::format("{}", arg->at(i));
                if (i < arg->size() - 1) {
                  vector_str += ", ";
                }
              }
              vector_str += ")";
              return vector_str;
            }
          },
          field);

      parts.push_back(std::move(part));
    }

    std::string line;
    if (!parts.empty()) {
      line = parts[0];
      for (size_t i = 1; i < parts.size(); ++i) {
        line += "|" + parts[i];
      }
    }

    fmt::println("{}", line);
  }
  fmt::println("");
}

enum class NodeType {
  HASH_JOIN,
  SCAN,
};

enum class DataType {
  INT32,
  INT64,
  FLOAT,
  DOUBLE,
  STRING,
  VECTOR,
};

struct ScanNode {
  size_t base_table_id;
};

struct JoinNode {
  size_t left;        // left index of plan node
  size_t right;       // right index of plan node
  size_t left_attr;   // left join index of table
  size_t right_attr;  // right join index of table
};

struct PlanNode {
  std::variant<ScanNode, JoinNode> data;
  std::vector<std::tuple<size_t, DataType>> output_attrs;

  PlanNode(std::variant<ScanNode, JoinNode> data,
           std::vector<std::tuple<size_t, DataType>> output_attrs)
      : data(std::move(data)), output_attrs(std::move(output_attrs)) {}
};

struct Plan {
  std::vector<PlanNode> nodes;
  std::vector<Table> inputs;
  // std::vector<Table>         tables;
  size_t root;

  size_t new_join_node(size_t left, size_t right, size_t left_attr,
                       size_t right_attr,
                       std::vector<std::tuple<size_t, DataType>> output_attrs) {
    JoinNode join{
        .left = left,
        .right = right,
        .left_attr = left_attr,
        .right_attr = right_attr,
    };
    auto ret = nodes.size();
    nodes.emplace_back(join, std::move(output_attrs));
    return ret;
  }

  size_t new_scan_node(size_t base_table_id,
                       std::vector<std::tuple<size_t, DataType>> output_attrs) {
    ScanNode scan{.base_table_id = base_table_id};
    auto ret = nodes.size();
    nodes.emplace_back(scan, std::move(output_attrs));
    return ret;
  }

  size_t new_input(Table&& input) {
    auto ret = inputs.size();
    inputs.emplace_back(std::move(input));
    return ret;
  }
};

////////////////////////////////////////////////////////////
//                        OPERATOR                        //
////////////////////////////////////////////////////////////

template <typename T>
using HT = std::unordered_map<T, std::vector<size_t>>;

template <typename T>
void build_hash_table(
    Table::iterator begin, Table::iterator end,
    std::pair<Table::iterator, Table::iterator> range,  // begin of total data
    std::shared_ptr<HT<T>> ht, size_t col_id, std::shared_ptr<std::mutex> mu) {
  auto init = range.first;
  for (auto it = begin; it != end; ++it) {
    std::visit(
        [&ht, &mu, idx = std::distance(init, it)](const auto& key) {
          using Tk = std::decay_t<decltype(key)>;
          if constexpr (std::is_same_v<Tk, T>) {
            mu->lock();
            if (auto itr = ht->find(key); itr == ht->end()) {
              ht->emplace(key, std::vector<size_t>(1, idx));  // init with idx
            } else {
              itr->second.push_back(idx);  // add idx
            }
            mu->unlock();
          } else if constexpr (not std::is_same_v<Tk, std::monostate>) {
            throw std::runtime_error("wrong type of field");
          }
        },
        (*it)[col_id]);
  }
}

template <typename T>
void probe_hash_table(
    Table::iterator begin, Table::iterator end,
    std::shared_ptr<Table> build_table, std::shared_ptr<HT<T>> ht,
    std::shared_ptr<Table> result, size_t probe_col_id,
    const std::vector<std::tuple<size_t, DataType>>& output_attrs,
    std::shared_ptr<std::mutex> mu) {
  for (auto it = begin; it != end; ++it) {
    std::visit(
        [&](const auto& key) {
          using Tk = std::decay_t<decltype(key)>;
          if constexpr (std::is_same_v<Tk, T>) {
            // find all row index
            std::vector<size_t> row_idx;
            if (auto itr = ht->find(key); itr != ht->end()) {
              row_idx.insert(row_idx.end(), itr->second.begin(),
                             itr->second.end());
            }

            for (auto record_idx : row_idx) {
              std::vector<Data> record;
              auto& build_record = build_table->at(record_idx);
              record.reserve(output_attrs.size());

              for (auto [col_idx, _] : output_attrs) {
                if (col_idx < build_record.size()) {
                  record.emplace_back(build_record[col_idx]);
                } else {
                  record.emplace_back((*it)[col_idx - build_record.size()]);
                }
              }

              mu->lock();
              result->emplace_back(std::move(record));
              mu->unlock();
            }

          } else if constexpr (not std::is_same_v<Tk, std::monostate>) {
            throw std::runtime_error("wrong type of field");
          }
        },
        (*it)[probe_col_id]);
  }
}

// scan, projection or filter
void scan_project_filter_table(
    Table::iterator begin, Table::iterator end, std::shared_ptr<Table> result,
    const std::vector<std::tuple<size_t, DataType>>& output_attrs,
    std::shared_ptr<std::mutex> mu) {
  for (auto it = begin; it != end; ++it) {
    std::vector<Data> record;
    for (auto& [col_idx, _] : output_attrs) {
      record.emplace_back((*it)[col_idx]);
    }
    mu->lock();
    result->emplace_back(std::move(record));
    mu->unlock();
  }
}

Table merge_table(std::vector<Table>&& results) {
  Table t;
  for (const auto& res : results) {
    t.insert(t.end(), res.begin(), res.end());
  }
  return t;
}

///////////////////////////////////////////////////////////
//                        EXECUTE                        //
///////////////////////////////////////////////////////////

std::shared_ptr<Table> generate_resource_group(ResourceGroup* rg, Plan& plan,
                                               size_t node_idx);

template <typename T>
std::shared_ptr<Table> hash_join_algorithm(
    ResourceGroup* rg, std::shared_ptr<Table> left,
    std::shared_ptr<Table> right, const JoinNode& join,
    const std::vector<std::tuple<size_t, DataType>>& output_attrs) {
  auto global_ht = std::make_shared<HT<T>>();
  auto mu = std::make_shared<std::mutex>();
  auto ret = std::make_shared<Table>();

  auto build = make_task_set(
      left, build_hash_table<T>, VOID_FINALIZE,
      std::pair<Table::iterator, Table::iterator>{nullptr, nullptr}, global_ht,
      join.left_attr, mu);

  mu = std::make_shared<std::mutex>();
  auto probe = make_task_set(right, probe_hash_table<T>, VOID_FINALIZE, left,
                             global_ht, ret, join.right_attr, output_attrs, mu);
  rg->add_task_set(std::move(build));
  rg->add_task_set(std::move(probe));
  return ret;
}

std::shared_ptr<Table> execute_hash_join(
    ResourceGroup* rg, Plan& plan, const JoinNode& join,
    const std::vector<std::tuple<size_t, DataType>>& output_attrs) {
  auto left_table = generate_resource_group(rg, plan, join.left);
  auto right_table = generate_resource_group(rg, plan, join.right);

  std::shared_ptr<Table> ret;

  switch (std::get<1>(plan.nodes[join.left].output_attrs[join.left_attr])) {
    case DataType::INT32:
      ret = hash_join_algorithm<int32_t>(rg, std::move(left_table),
                                         std::move(right_table), join,
                                         output_attrs);
      break;
    case DataType::INT64:
      ret = hash_join_algorithm<int64_t>(rg, std::move(left_table),
                                         std::move(right_table), join,
                                         output_attrs);
      break;
    case DataType::FLOAT:
      ret = hash_join_algorithm<float>(rg, std::move(left_table),
                                       std::move(right_table), join,
                                       output_attrs);
      break;
    case DataType::DOUBLE:
      ret = hash_join_algorithm<double>(rg, std::move(left_table),
                                        std::move(right_table), join,
                                        output_attrs);
      break;
    case DataType::STRING:
      ret = hash_join_algorithm<std::string>(rg, std::move(left_table),
                                             std::move(right_table), join,
                                             output_attrs);
      break;
  }

  return ret;
}

std::shared_ptr<Table> execute_scan(
    ResourceGroup* rg, Plan& plan, const ScanNode& scan,
    const std::vector<std::tuple<size_t, DataType>>& output_attrs) {
  auto table_id = scan.base_table_id;
  auto& input = plan.inputs[table_id];

  auto mu = std::make_shared<std::mutex>();
  auto ret = std::make_shared<Table>();

  auto taskset =
      make_task_set(input.begin(), input.end(), scan_project_filter_table,
                    VOID_FINALIZE, ret, output_attrs, mu);
  rg->add_task_set(std::move(taskset));

  return ret;
}

std::shared_ptr<Table> generate_resource_group(ResourceGroup* rg, Plan& plan,
                                               size_t node_idx) {
  auto& node = plan.nodes[node_idx];

  return std::visit(
      [&](const auto& value) {
        using T = std::decay_t<decltype(value)>;
        if constexpr (std::is_same_v<T, JoinNode>) {
          return execute_hash_join(rg, plan, value, node.output_attrs);
        } else {
          return execute_scan(rg, plan, value, node.output_attrs);
        }
      },
      node.data);
}

}  // namespace hashjoin
// NOLINTEND
